import codecs
import sys

base = '/home/rickwwang/project_research/fire_data/running/'
# prompt_rank_file = base + 'conv-s2s-d0.3_bpe500/40.score.txt.prank'
# coherence_rank_file = base + 'conv-s2s-d0.3_bpe500/40.score.txt.coherence'
# print ('conv')

# prompt_rank_file = base + 'conv-s2s-d0.3-fusion_bpe500/18.score.txt.prank'
# coherence_rank_file = base + 'conv-s2s-d0.3-fusion_bpe500/18.score.txt.coherence'
# print ('conv-fusion')

# prompt_rank_file = base + 'transformer-lm_bpe500_warmlr_format2/40.score.txt.prank'
# coherence_rank_file = base + 'transformer-lm_bpe500_warmlr_format2/40.score.txt.coherence'
# print ('trans-lm')

# prompt_rank_file = base + 'transformer-lm_bpe500_warmlr_format_pre2/27.score.txt.prank'
# coherence_rank_file = base + 'transformer-lm_bpe500_warmlr_format_pre2/27.score.txt.coherence'
# print ('trans-lm-pre')

mode = sys.argv[1]
prompt_rank_file = sys.argv[2]
coherence_rank_file = sys.argv[3]
# print ('trans-lm-pre')


def get_score(line):
    elements = line.strip().split('\t')
    return int(elements[0]), float(elements[2])


if mode in ['p', 'pc']:
    # prompt rank
    p_scores = codecs.open(prompt_rank_file, 'r', encoding='utf8').readlines()
    p_scores = map(get_score, p_scores)
    prompt_id_to_score = dict(p_scores)

    total_story_pr = 10000
    prompt_each_story = 10

    # prompt accuracy
    count = 0
    correct_count = 0
    ids = range(total_story_pr)
    seg_ids = [ids[i:i + prompt_each_story] for i in range(0, len(ids), prompt_each_story)]
    for seg in seg_ids:
        seg_scores = [prompt_id_to_score[s] for s in seg if s in prompt_id_to_score]
        if len(seg_scores) == prompt_each_story:
            count += 1
            if seg_scores[0] == max(seg_scores):
                correct_count += 1
    print ('total count:', count)
    print ('prompt accuracy:', correct_count * 1.0 / count)

if mode in ['c', 'pc']:
    # coherence rank
    c_scores = codecs.open(coherence_rank_file, 'r', encoding='utf8').readlines()
    c_scores = map(get_score, c_scores)
    coherence_id_to_score = dict(c_scores)

    total_story_cr = 197100
    shuffle_each_story = 15

    # coherence accuracy
    count = 0
    correct_count = 0
    ids = range(total_story_cr)
    seg_ids = [ids[i:i + shuffle_each_story] for i in range(0, len(ids), shuffle_each_story)]
    for seg in seg_ids:
        seg_scores = [coherence_id_to_score[s] for s in seg if s in coherence_id_to_score]
        if len(seg_scores) == shuffle_each_story:
            count += 1
            if seg_scores[0] == max(seg_scores):
                correct_count += 1
    print ('total count:', count)
    print ('coherence accuracy:', correct_count * 1.0 / count)
